﻿using System;
using System.ComponentModel;
using System.Runtime.InteropServices;
using System.Text;

namespace Kobets.Automation.Chimera.Common
{
    /// <summary>
    /// Helper class for Win32 API
    /// </summary>
    public static class WinApi
    {
        #region Vars

        internal const int HwndBroadcast = 0xffff;
        internal const int SwShownormal = 1;

        public const int ReadControl = 0x00020000;
        public const int StandardRightsRequired = 0x000F0000;
        public const int StandardRightsRead = ReadControl;
        public const int StandardRightsWrite = ReadControl;
        public const int StandardRightsExecute = ReadControl;
        public const int StandardRightsAll = 0x001F0000;
        public const int SpecificRightsAll = 0x0000FFFF;
        public const int TokenAssignPrimary = 0x0001;
        public const int TokenDuplicate = 0x0002;
        public const int TokenImpersonate = 0x0004;
        public const int TokenQuery = 0x0008;
        public const int TokenQuerySource = 0x0010;
        public const int TokenAdjustPrivileges = 0x0020;
        public const int TokenAdjustGroups = 0x0040;
        public const int TokenAdjustDefault = 0x0080;
        public const int TokenAdjustSessionid = 0x0100;
        public const int TokenAllAccessP = (StandardRightsRequired | TokenAssignPrimary | TokenDuplicate | TokenImpersonate | TokenQuery | TokenQuerySource | TokenAdjustPrivileges | TokenAdjustGroups | TokenAdjustDefault);
        public const int TokenAllAccess = TokenAllAccessP | TokenAdjustSessionid;
        public const int TokenRead = StandardRightsRead | TokenQuery;
        public const int TokenWrite = StandardRightsWrite | TokenAdjustPrivileges | TokenAdjustGroups | TokenAdjustDefault;
        public const int TokenExecute = StandardRightsExecute;
        public const uint MaximumAllowed = 0x2000000;
        public const int CreateNewProcessGroup = 0x00000200;
        public const int CreateUnicodeEnvironment = 0x00000400;
        public const int IdlePriorityClass = 0x40;
        public const int NormalPriorityClass = 0x20;
        public const int HighPriorityClass = 0x80;
        public const int RealtimePriorityClass = 0x100;
        public const int CreateNewConsole = 0x00000010;
        public const string SeDebugName = "SeDebugPrivilege";
        public const string SeRestoreName = "SeRestorePrivilege";
        public const string SeBackupName = "SeBackupPrivilege";
        public const int SePrivilegeEnabled = 0x0002;
        public const int ErrorNotAllAssigned = 1300;
        
        public static int InvalidHandleValue = -1;

        [StructLayout(LayoutKind.Sequential)]
        public struct SecurityAttributes
        {
            public int Length;
            public IntPtr lpSecurityDescriptor;
            public bool bInheritHandle;
        }

        public enum TokenInformationClass
        {
            TokenUser = 1,
            TokenGroups,
            TokenPrivileges,
            TokenOwner,
            TokenPrimaryGroup,
            TokenDefaultDacl,
            TokenSource,
            TokenType,
            TokenImpersonationLevel,
            TokenStatistics,
            TokenRestrictedSids,
            TokenSessionId,
            TokenGroupsAndPrivileges,
            TokenSessionReference,
            TokenSandBoxInert,
            TokenAuditPolicy,
            TokenOrigin,
            MaxTokenInfoClass
        }

        [StructLayout(LayoutKind.Sequential)]
        public struct Startupinfo
        {
            public int cb;
            public String lpReserved;
            public String lpDesktop;
            public String lpTitle;
            public uint dwX;
            public uint dwY;
            public uint dwXSize;
            public uint dwYSize;
            public uint dwXCountChars;
            public uint dwYCountChars;
            public uint dwFillAttribute;
            public uint dwFlags;
            public short wShowWindow;
            public short cbReserved2;
            public IntPtr lpReserved2;
            public IntPtr hStdInput;
            public IntPtr hStdOutput;
            public IntPtr hStdError;
        }

        [StructLayout(LayoutKind.Sequential)]
        public struct ProcessInformation
        {
            public IntPtr hProcess;
            public IntPtr hThread;
            public uint dwProcessId;
            public uint dwThreadId;
        }

        [StructLayout(LayoutKind.Sequential)]
        public struct Luid
        {
            public int LowPart;
            public int HighPart;
        }

        [StructLayout(LayoutKind.Sequential)]
        internal struct LuidAndAtributes
        {
            public Luid Luid;
            public int Attributes;
        }

        [StructLayout(LayoutKind.Sequential)]
        public struct TokenPrivileges
        {
            internal int PrivilegeCount;
            [MarshalAs(UnmanagedType.ByValArray, SizeConst = 3)]
            internal int[] Privileges;
        }

        [StructLayout(LayoutKind.Sequential)]
        public struct Processentry32
        {
            public uint dwSize;
            public uint cntUsage;
            public uint th32ProcessID;
            public IntPtr th32DefaultHeapID;
            public uint th32ModuleID;
            public uint cntThreads;
            public uint th32ParentProcessID;
            public int pcPriClassBase;
            public uint dwFlags;
            [MarshalAs(UnmanagedType.ByValTStr, SizeConst = 260)]
            public string szExeFile;
        }

        #endregion

        #region Externs

        [DllImport("user32")]
        public static extern int RegisterWindowMessage(string message);

        [DllImport("user32")]
        public static extern bool PostMessage(IntPtr hwnd, int msg, IntPtr wparam, IntPtr lparam);

        [DllImport("user32.dll")]
        public static extern bool ShowWindow(IntPtr hWnd, int nCmdShow);

        [DllImport("user32.dll")]
        public static extern bool SetForegroundWindow(IntPtr hWnd);

        [DllImport("kernel32.dll")]
        public static extern int WTSGetActiveConsoleSessionId();

        [DllImport("Wtsapi32.dll")]
        public static extern bool WTSQueryUserToken(uint sessionId, ref IntPtr phToken);

        [DllImport("advapi32.dll", SetLastError = true)]
        public static extern bool LookupPrivilegeValue(IntPtr lpSystemName, string lpname, [MarshalAs(UnmanagedType.Struct)] ref Luid lpLuid);

        [DllImport("advapi32.dll", EntryPoint = "CreateProcessAsUser", SetLastError = true, CharSet = CharSet.Ansi, CallingConvention = CallingConvention.StdCall)]
        public extern static bool CreateProcessAsUser(IntPtr hToken, string lpApplicationName, string lpCommandLine, IntPtr lpProcessAttributes, IntPtr lpThreadAttributes, bool bInheritHandle, int dwCreationFlags, IntPtr lpEnvironment, string lpCurrentDirectory, ref Startupinfo lpStartupInfo, out ProcessInformation lpProcessInformation);

        [DllImport("advapi32.dll", CharSet = CharSet.Auto, SetLastError = true)]
        public extern static bool DuplicateToken(IntPtr existingTokenHandle, int securityImpersonationLevel, ref IntPtr duplicateTokenHandle);

        [DllImport("advapi32.dll", EntryPoint = "DuplicateTokenEx")]
        public extern static bool DuplicateTokenEx(IntPtr existingTokenHandle, uint dwDesiredAccess, ref SecurityAttributes lpThreadAttributes, int tokenType, int impersonationLevel, ref IntPtr duplicateTokenHandle);

        [DllImport("advapi32.dll", SetLastError = true)]
        public static extern bool AdjustTokenPrivileges(IntPtr tokenHandle, bool disableAllPrivileges, ref TokenPrivileges newState, int bufferLength, IntPtr previousState, IntPtr returnLength);

        [DllImport("advapi32.dll", SetLastError = true)]
        public static extern bool SetTokenInformation(IntPtr tokenHandle, TokenInformationClass tokenInformationClass, uint tokenInformation, uint tokenInformationLength);

        [DllImport("userenv.dll", SetLastError = true)]
        public static extern bool CreateEnvironmentBlock(out IntPtr lpEnvironment, IntPtr hToken, bool bInherit);

        [DllImport("userenv.dll", SetLastError = true)]
        public static extern bool DestroyEnvironmentBlock(IntPtr lpEnvironment);

        #endregion

        public static int RegisterWindowMessage(string format, params object[] args)
        {
            var message = String.Format(format, args);
            return RegisterWindowMessage(message);
        }

        /// <summary>
        /// Brings window to front, even if its minimized or in tray
        /// </summary>
        /// <param name="window"></param>
        public static void ShowToFront(IntPtr window)
        {
            ShowWindow(window, SwShownormal);
            SetForegroundWindow(window);
        }

        /// <summary>
        /// Runs a process via impersonation
        /// </summary>
        /// <param name="filename"></param>
        /// <param name="commandLine"></param>
        /// <param name="workingDirectory"></param>
        /// <param name="sessionId">if equals -1 (default) will run in console/admin session</param>
        /// <returns></returns>
        public static int RunAs(string filename, string commandLine, string workingDirectory, int sessionId = -1)
        {
            var hToken = new IntPtr();
            var si = new Startupinfo();

            var dwSessionId = sessionId == -1 ? WTSGetActiveConsoleSessionId() : sessionId;

            if (WTSQueryUserToken((uint)dwSessionId, ref hToken))
            {
                IntPtr pEnvBlock;

                if (CreateEnvironmentBlock(out pEnvBlock, hToken, false))
                {
                    var cmdLine = new StringBuilder(commandLine, 32768);
                    ProcessInformation pi;

                    if (CreateProcessAsUser(hToken, filename, cmdLine.ToString(), IntPtr.Zero, IntPtr.Zero, false, CreateUnicodeEnvironment, pEnvBlock, workingDirectory, ref si, out pi))
                    {
                        // sing hallelluya
                        return (int)pi.dwProcessId;
                    }
                    
                    throw new Win32Exception();
                }

                throw new Win32Exception();
            }

            throw new Win32Exception();
        }
    }
}
